package main.java.LeetCode.ClassicalAlgorithm.Graph.Impl;

import java.util.HashMap;

public class UnionFind<E> {

    private final HashMap<E,E> parents;
    private final HashMap<E,Integer> sizes;


    public UnionFind() {
        parents = new HashMap<>();
        sizes = new HashMap<>();
    }

    public void union(E p,E q){
        init(p);
        init(q);

        E parentP = findParent(p,0);
        E parentQ = findParent(q,0);

        if(sizes.get(parentP)>sizes.get(parentQ)){
            parents.put(parentQ,parentP);
            sizes.put(parentP,sizes.get(parentP)+sizes.get(parentQ));
        }else {
            parents.put(parentP,parentQ);
            sizes.put(parentQ,sizes.get(parentP)+sizes.get(parentQ));
        }
    }
    public boolean isConnected(E p,E q){
        if(!parents.containsKey(p)||!parents.containsKey(q))
            return false;
        return findParent(p,0).equals(findParent(q,0));
    }

    private E findParent(E p,int count){
        if(!parents.get(p).equals(p)){
            parents.put(p,findParent(parents.get(p),count+1));
            sizes.put(p,sizes.get(p) - count);
        }
        return parents.get(p);
    }

    public E getParent(E p){
        return findParent(p,0);
    }

    private void init(E p){
        if(!parents.containsKey(p)){
            parents.put(p,p);
            sizes.put(p,1);
        }
    }
}
